from typing import Any, Dict, List, Optional, Tuple

from pydantic import Field, field_validator

from src.agent.toolcall import ToolCallAgent
from src.config import config
from src.logger import logger
from src.prompt.mcp import MULTIMEDIA_RESPONSE_PROMPT, NEXT_STEP_PROMPT, SYSTEM_PROMPT
from src.schema import Message
from src.tool.base import ToolResult


class MCPAgent(ToolCallAgent):
    """MCP (Model Context Protocol) 协议代理。

    这个代理专门用于与 MCP 服务器交互，支持通过多种传输方式连接：
    1. SSE (Server-Sent Events) 传输：适用于 HTTP/WebSocket 连接
    2. Stdio 传输：适用于本地进程通信
    
    主要功能：
    - 动态连接和断开 MCP 服务器
    - 自动发现和刷新服务器提供的工具
    - 支持多服务器并发连接
    - 处理连接失败和恢复
    - 工具模式变更检测和通知
    - 多媒体内容支持
    
    适用场景：
    - 需要使用外部 MCP 服务提供的工具
    - 动态工具集合的任务
    - 分布式工具调用
    - 插件化架构的应用
    """

    # 代理基本信息
    name: str = Field(default="mcp_agent", description="MCP 代理名称标识")
    description: str = Field(
        default="连接到 MCP 服务器并使用其工具的智能代理",
        description="代理功能描述"
    )
    system_prompt: str = Field(
        default=SYSTEM_PROMPT, description="MCP 代理系统提示词"
    )
    next_step_prompt: str = Field(
        default=NEXT_STEP_PROMPT, description="MCP 代理步骤指导提示词"
    )

    # 执行配置参数
    max_steps: int = Field(default=10, description="最大执行步数")
    max_observe: int = Field(
        default=10000, description="工具输出观察的最大字符数"
    )
    refresh_tools_interval: int = Field(
        default=5, description="每 N 步刷新一次工具列表"
    )

    # MCP 客户端和工具管理
    mcp_clients: Any = Field(
        default=None, description="MCP 客户端实例，用于与服务器通信"
    )
    available_tools: Optional[Any] = Field(
        default=None, description="可用工具集合，在初始化时设置"
    )

    # 状态跟踪和管理
    tool_schemas: Dict[str, Dict[str, Any]] = Field(
        default_factory=dict, description="工具模式定义缓存，用于检测变更"
    )
    special_tool_names: List[str] = Field(
        default=["terminate"], description="特殊工具名称列表"
    )
    connected_servers: Dict[str, str] = Field(
        default_factory=dict, description="已连接的服务器映射：server_id -> url/command"
    )
    initialized: bool = Field(
        default=False, description="初始化状态标记"
    )

    @classmethod
    @field_validator("mcp_clients", mode="before")
    def init_mcp_clients(cls, v: Any) -> Any:
        """初始化 MCP 客户端（如果未提供）。
        
        这个验证器确保 MCP 代理始终有一个有效的 MCPClients 实例。
        如果用户没有提供，则创建默认实例。
        
        Args:
            v: 用户提供的 mcp_clients 值
            
        Returns:
            MCPClients: 有效的 MCP 客户端实例
        """
        if v is None:
            from src.tool.mcp_client import MCPClients
            return MCPClients()
        return v

    @classmethod
    async def create(cls, **kwargs) -> "MCPAgent":
        """工厂方法：创建并正确初始化 MCPAgent 实例。
        
        这个工厂方法提供了安全的 MCP 代理创建方式，确保在返回实例之前
        所有必要的初始化步骤都正确完成。
        
        Args:
            **kwargs: 代理构造参数
            
        Returns:
            MCPAgent: 完全初始化的 MCP 代理实例
        """
        # 创建代理实例
        instance = cls(**kwargs)
        try:
            # 尝试初始化 MCP 服务器连接
            await instance.initialize_mcp_servers()
        except Exception as e:
            # 如果初始化失败，记录警告但不中断
            logger.warning(
                f"初始化 MCP 服务器失败: {str(e)}。继续使用本地工具。"
            )

        # 标记为已初始化
        instance.initialized = True

        # 设置可用工具集合（如果未设置）
        if instance.available_tools is None:
            instance.available_tools = instance.mcp_clients

        # 如果有 MCP 客户端，添加可用工具的系统消息
        if instance.mcp_clients and hasattr(instance.mcp_clients, "tool_map"):
            tool_names = list(instance.mcp_clients.tool_map.keys())
            instance.memory.add_message(
                Message.system_message(
                    f"{instance.system_prompt}\n\n可用的 MCP 工具: {', '.join(tool_names)}"
                )
            )

        return instance

    async def initialize_mcp_servers(self) -> None:
        """初始化到配置的 MCP 服务器的连接。
        
        这个方法根据配置文件中的设置，逐个初始化所有 MCP 服务器连接。
        支持 SSE 和 Stdio 两种连接类型。
        
        如果某个服务器连接失败，会记录错误但不会中断整个初始化过程。
        """
        # 遍历配置中的所有 MCP 服务器
        for server_id, server_config in config.mcp_config.servers.items():
            try:
                # 根据连接类型选择相应的连接方法
                if server_config.type == "sse" and server_config.url:
                    # SSE (Server-Sent Events) 连接
                    await self.connect_mcp_server(server_config.url, server_id)
                    logger.info(
                        f"已连接到 MCP 服务器 {server_id}，地址: {server_config.url}"
                    )
                elif server_config.type == "stdio" and server_config.command:
                    # Stdio (标准输入输出) 连接
                    await self.connect_mcp_server(
                        server_config.command,
                        server_id,
                        use_stdio=True,
                        stdio_args=server_config.args,
                    )
                    logger.info(
                        f"已连接到 MCP 服务器 {server_id}，使用命令: {server_config.command}"
                    )
            except Exception as e:
                # 记录连接失败但继续尝试其他服务器
                logger.error(f"连接到 MCP 服务器 {server_id} 失败: {e}")

    async def connect_mcp_server(
        self,
        server_url: str,
        server_id: str = "",
        use_stdio: bool = False,
        stdio_args: Optional[List[str]] = None,
    ) -> None:
        """连接到 MCP 服务器并添加其工具。
        
        这个方法负责建立与单个 MCP 服务器的连接，并将服务器提供的
        工具添加到代理的可用工具集合中。
        
        Args:
            server_url: 服务器 URL 或命令路径
            server_id: 服务器唯一标识符
            use_stdio: 是否使用 stdio 传输
            stdio_args: stdio 连接的额外参数
        """
        # 根据传输类型选择连接方法
        if use_stdio:
            # 使用 stdio 连接（本地进程）
            await self.mcp_clients.connect_stdio(
                server_url, stdio_args or [], server_id
            )
        else:
            # 使用 SSE 连接（HTTP/WebSocket）
            await self.mcp_clients.connect_sse(server_url, server_id)

        # 记录已连接的服务器
        self.connected_servers[server_id or server_url] = server_url

        # 如果需要，将新连接的服务器工具添加到可用工具集合
        if self.available_tools and self.available_tools is not self.mcp_clients:
            # 筛选属于当前服务器的工具
            new_tools = [
                tool for tool in self.mcp_clients.tools if tool.server_id == server_id
            ]
            self.available_tools.add_tools(*new_tools)

        # 刷新工具模式定义
        await self._refresh_tools()

    async def disconnect_mcp_server(self, server_id: str = "") -> None:
        """断开与 MCP 服务器的连接并移除其工具。
        
        这个方法负责安全地断开与 MCP 服务器的连接，并从可用工具集合中
        移除该服务器提供的所有工具。
        
        Args:
            server_id: 要断开的服务器 ID，空字符串表示断开所有服务器
        """
        # 检查 MCP 客户端是否存在
        if not self.mcp_clients:
            return

        # 断开指定服务器或所有服务器
        await self.mcp_clients.disconnect(server_id)

        # 更新已连接服务器列表
        if server_id:
            # 移除指定服务器
            self.connected_servers.pop(server_id, None)
        else:
            # 清空所有连接
            self.connected_servers.clear()

        # Rebuild available_tools if needed
        if self.available_tools and self.available_tools is not self.mcp_clients:
            from src.tool.mcp_client import MCPClientTool

            base_tools = [
                tool
                for tool in self.available_tools.tools
                if not isinstance(tool, MCPClientTool)
            ]

            remaining_mcp_tools = [
                tool
                for tool in self.mcp_clients.tools
                if hasattr(tool, "server_id")
                and tool.server_id in self.connected_servers
            ]

            self.available_tools.tools = base_tools
            self.available_tools.add_tools(*remaining_mcp_tools)

    async def initialize(
        self,
        connection_type: Optional[str] = None,
        server_url: Optional[str] = None,
        command: Optional[str] = None,
        args: Optional[List[str]] = None,
        server_id: str = "",
    ) -> None:
        """Initialize a single MCP connection."""
        # Allow initialization without connection parameters (use configured servers)
        if connection_type is None:
            if not self.initialized:
                await self.initialize_mcp_servers()
        elif connection_type == "sse":
            if not server_url:
                raise ValueError("Server URL is required for SSE connection")
            await self.connect_mcp_server(server_url=server_url, server_id=server_id)
        elif connection_type == "stdio":
            if not command:
                raise ValueError("Command is required for stdio connection")
            await self.connect_mcp_server(
                server_url=command,
                server_id=server_id,
                use_stdio=True,
                stdio_args=args or [],
            )
        else:
            raise ValueError(f"Unsupported connection type: {connection_type}")

        # Set available_tools if needed
        if self.available_tools is None:
            self.available_tools = self.mcp_clients

        await self._refresh_tools()
        self.initialized = True

    async def _refresh_tools(self) -> Tuple[List[str], List[str]]:
        """从 MCP 服务器刷新可用工具列表。
        
        这个方法定期检查 MCP 服务器上可用的工具，检测工具的增加、删除或修改。
        当检测到变化时，会自动更新代理的对话历史以通知用户。
        
        Returns:
            Tuple[List[str], List[str]]: (新增工具列表, 移除工具列表)
        """
        # 检查 MCP 客户端和会话是否存在
        if not self.mcp_clients or not self.mcp_clients.sessions:
            return [], []

        # 获取当前工具模式定义
        response = await self.mcp_clients.sessions.list_tools()
        current_tools = {tool.name: tool.inputSchema for tool in response.tools}

        # 确定工具变更
        current_names = set(current_tools.keys())
        previous_names = set(self.tool_schemas.keys())

        # 计算工具变更
        added_tools = list(current_names - previous_names)  # 新增工具
        removed_tools = list(previous_names - current_names)  # 移除工具

        # 检查工具模式定义的变化
        changed_tools = [
            name
            for name in current_names.intersection(previous_names)
            if current_tools[name] != self.tool_schemas.get(name)
        ]

        # 更新存储的工具模式
        self.tool_schemas = current_tools

        # 记录并通知工具变化
        if added_tools:
            logger.info(f"新增 MCP 工具: {added_tools}")
            self.memory.add_message(
                Message.system_message(f"新可用工具: {', '.join(added_tools)}")
            )

        if removed_tools:
            logger.info(f"移除 MCP 工具: {removed_tools}")
            self.memory.add_message(
                Message.system_message(
                    f"不再可用的工具: {', '.join(removed_tools)}"
                )
            )

        if changed_tools:
            logger.info(f"变更 MCP 工具: {changed_tools}")

        return added_tools, removed_tools

    async def think(self) -> bool:
        """处理当前状态并决定下一步操作。
        
        这个方法重写了父类的 think 方法，增加了 MCP 特有的逻辑：
        1. 确保 MCP 服务器已初始化
        2. 定期刷新工具列表
        3. 检测服务器状态
        
        Returns:
            bool: True 表示需要执行行动，False 表示只需返回思考结果
        """
        # 如果未初始化，先初始化 MCP 服务器
        if not self.initialized:
            await self.initialize_mcp_servers()
            self.initialized = True

        # 定期刷新工具（如果 MCP 客户端存在）
        if (
            self.mcp_clients
            and self.current_step % self.refresh_tools_interval == 0
            and self.mcp_clients.sessions
        ):
            await self._refresh_tools()
            # 检测 MCP 服务是否已关闭
            if self.mcp_clients.sessions and not self.mcp_clients.tool_map:
                logger.info("MCP 服务已关闭")

        # 调用父类的 think 方法执行核心思考逻辑
        return await super().think()

    async def _handle_special_tool(self, name: str, result: Any, **kwargs) -> None:
        """处理特殊工具的执行和状态变化。
        
        MCP 代理的特殊工具处理，在父类功能基础上增加了多媒体内容处理。
        
        Args:
            name: 工具名称
            result: 工具执行结果
            **kwargs: 额外参数
        """
        # 调用父类的特殊工具处理
        await super()._handle_special_tool(name, result, **kwargs)

        # 处理多媒体响应（如图像）
        if isinstance(result, ToolResult) and result.base64_image:
            self.memory.add_message(
                Message.system_message(
                    MULTIMEDIA_RESPONSE_PROMPT.format(tool_name=name)
                )
            )

    def _should_finish_execution(self, name: str, **kwargs) -> bool:
        """判断工具执行是否应该终止代理。
        
        MCP 代理只有在遇到 'terminate' 工具时才会终止执行。
        
        Args:
            name: 工具名称
            **kwargs: 额外参数
            
        Returns:
            bool: True 表示应该终止代理执行
        """
        return name.lower() == "terminate"

    async def cleanup(self) -> None:
        """完成时清理 MCP 连接。
        
        这个方法重写了父类的 cleanup 方法，专门处理 MCP 相关的资源清理。
        主要是断开所有 MCP 服务器连接并重置初始化状态。
        """
        # 如果已初始化且有 MCP 客户端，断开所有连接
        if self.initialized and self.mcp_clients:
            await self.disconnect_mcp_server()  # 断开所有服务器
            self.initialized = False  # 重置初始化状态
            logger.info("已关闭 MCP 连接")
